組裝 Transformer

隨著這幾天的逐步逐步的實做，目前一個 Transformer 模組中要有的元素都已經有了，下一步是要將許多 Transformer 組裝起來變成一個 Transformer Block：

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.att = MultiHeadAttention(
            d_in=cfg["emb_dim"],
            d_out=cfg["emb_dim"],
            context_length=cfg["context_length"],
            num_heads=cfg["n_heads"], 
            dropout=cfg["drop_rate"],
            qkv_bias=cfg["qkv_bias"])
        self.ff = FeedForward(cfg)
        self.norm1 = LayerNorm(cfg["emb_dim"])
        self.norm2 = LayerNorm(cfg["emb_dim"])
        self.drop_shortcut = nn.Dropout(cfg["drop_rate"])

    def forward(self, x):
        shortcut = x
        x = self.norm1(x)
        x = self.att(x)  
        x = self.drop_shortcut(x)
        x = x + shortcut 

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = self.drop_shortcut(x)
        x = x + shortcut 

        return x

完整的順序如下，input 是 tokenize 並經過 Embedding 的矩陣，接著先經過一次 normalization layer，接著透過 multi head attention module 並且 dropout 出結果，同時在 normalization 前拉出一個 shortcut connection 連結。

這一段的結果會再經過一次 normalization 送進 feed forward 中，feed forward 又包括了 linear layer 與 GELU activation 後，再送出 dropout 第一次，同樣的 input & output 會透過 shortcut connection 連結。

這一整塊便是 Transformer 模組，綜觀來看，從最一開始經過 attention 讓我們可以掌握文字之間關連，到後續 feed forward 單純去預測什麼是最佳文字結果。

而過程中即便維度會轉置、重新編碼計算，但始終保持最終輸出結果不變也是其中一個輸出重點。

組裝 GPT

最終經過 transformer 後，會再做一次最後最後一次的 normalization 最後透過 linear layer 輸出完整結果。

class GPTModel(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.pos_emb = nn.Embedding(cfg["context_length"], cfg["emb_dim"])
        self.drop_emb = nn.Dropout(cfg["drop_rate"])
        
        self.trf_blocks = nn.Sequential(
            *[TransformerBlock(cfg) for _ in range(cfg["n_layers"])])
        
        self.final_norm = LayerNorm(cfg["emb_dim"])
        self.out_head = nn.Linear(
            cfg["emb_dim"], cfg["vocab_size"], bias=False
        )

    def forward(self, in_idx):
        batch_size, seq_len = in_idx.shape
        tok_embeds = self.tok_emb(in_idx)
        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))
        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]
        x = self.drop_emb(x)
        x = self.trf_blocks(x)
        x = self.final_norm(x)
        logits = self.out_head(x)
        return logits

於是我們再將 Transformer Block 加入進完整的 GPTModel Class 之中。接著再加入生成文字的 function:

def generate_text_simple(model, idx, max_new_tokens, context_size):
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -context_size:]
        
        with torch.no_grad():
            logits = model(idx_cond)
        
        logits = logits[:, -1, :]  
        probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)
        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)
        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)

    return idx

這個 function 會負責讀取一段文字，生成下一個文字後加入在文字之後，接著繼續生成下一個。而也因為我們還沒訓練模型，所以先關閉 dropout 等等的模組。

最終就可以做到一個胡言亂語生成文字的 GPT 模組了。